import torch
from torch_geometric.data import InMemoryDataset
from tqdm import tqdm
from sklearn.preprocessing import LabelEncoder
import pandas as pd
from torch_geometric.data import Data
import numpy as np
#数据读取 获取点击数据 因为没有列名所以header=none
df = pd.read_csv('yoochoose-clicks.dat', header=None)
#为DataFrame对象的列设置名称。'session_id'，'timestamp'，'item_id'和'category'是新设置的列名。
df.columns = ['session_id', 'timestamp', 'item_id', 'category']
#数据读取获得购买数据
buy_df = pd.read_csv('yoochoose-buys.dat', header=None)
buy_df.columns = ['session_id', 'timestamp', 'item_id', 'price', 'quantity']

#数据有点多1个多G，咱们只选择其中一小部分来建模


#df.session_id.unique()从数据中挑选100000个不同的sessionid  replace=False每个序号只能选一次
sampled_session_id = np.random.choice(df.session_id.unique(), 100000, replace=False)
#从原来一个多G的数据中挑选出和上一步10万个sessionid相关的数据
#df.session_id.isin(sampled_session_id)判断sessionid是否在原数据中
df = df.loc[df.session_id.isin(sampled_session_id)]
#label为用户是否购买 在df浏览表中查询buy_df购买表 如果有则为true代表购买该商品
df['label'] = df.session_id.isin(buy_df.session_id)
#数据读取类
class YooChooseBinaryDataset(InMemoryDataset):
    #不用动
    def __init__(self, root, transform=None, pre_transform=None):
        super(YooChooseBinaryDataset, self).__init__(root, transform, pre_transform)  # transform就是数据增强，对每一个数据都执行
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):  # 检查self.raw_dir目录下是否存在raw_file_names()属性方法返回的每个文件
        # 如有文件不存在，则调用download()方法执行原始文件下载
        return []

    @property
    def processed_file_names(self):  # 检查self.processed_dir目录下是否存在self.processed_file_names属性方法返回的所有文件，没有就会走process
        return ['yoochoose_click_binary_1M_sess.dataset']

    def download(self):
        pass

    def process(self):
        data_list = []

        # process by session_id
        #以session_id为组划分 grouped是一个groupby对象
        grouped = df.groupby('session_id')
        #每个grouped对象包含一个session_id和一个与session_id相关的dataframe
        for session_id, group in tqdm(grouped):
            # 把itemid变成从0开始的数组 例原始数据为[39902,39902,57,57]转为[1,1,0,0]
            sess_item_id = LabelEncoder().fit_transform(group.item_id)
            #group.reset_index用于重新设置dataframe的索引 drop=True表示再重置索引的时候删除原来的索引
            group = group.reset_index(drop=True)
            #添加新的索引'sess_item_id'即用户新的item_id 例从[39902,39902,57,57]转为[1,1,0,0]
            group['sess_item_id'] = sess_item_id
            #group.loc用于选择group.session_id == session_id的行按sess_item_id排序并返回'sess_item_id', 'item_id'数据 item_id.drop_duplicates().values删除重复的itemid
            node_features = group.loc[group.session_id == session_id, ['sess_item_id', 'item_id']].sort_values(
                'sess_item_id').item_id.drop_duplicates().values
            #把node_features转为tensor类型 node features存的是item_id
            #unsqueeze增加维度 参数为下标从0开始 即从原来的(2,)变成([2, 1])在第二个下标加一个维度
            print("node_features", node_features.shape)
            #假设之前的sess_item_id为[0,1,2,3] 因为我们初始为按时间排序 所以顺序为0->1->2->3我们想要建设多个二维向量来保存这个数据
            """
            group.sess_item_id.values[:-1]即[0,1,2]
            group.sess_item_id.values[1:]即[1,2,3]
            torch.tensor([source_nodes, target_nodes]合并两个向量
            每一列为一对连接信息 例第一列为[0]
                                      [1]表示从0到1
            [[0,1,2] 
             [1,2,3]]
            """
            target_nodes = group.sess_item_id.values[1:]
            source_nodes = group.sess_item_id.values[:-1]
            edge_index = torch.tensor([source_nodes, target_nodes], dtype=torch.long)
            #x为特征向量
            x = node_features
            #y为标签
            y = torch.FloatTensor([group.label.values[0]])
            #edge_index表示图的关系
            data = Data(x=x, edge_index=edge_index, y=y)
            #一个data代表一个seesionid所有数据即一个用户一次登录的行为 即一个图 datalist里面包含多张图
            data_list.append(data)
        #data保存了所有数据 slices保存切片信息比如位置方便以后小批量minibatch训练 一个slice就是一个batch
        data, slices = self.collate(data_list)
        #保存到指定位置即./data/processed中
        torch.save((data, slices), self.processed_paths[0])

#调用数据加载存到当前目录的data下
dataset = YooChooseBinaryDataset(root='data/')

#设置嵌入的向量
embed_dim = 128
from torch_geometric.nn import TopKPooling, SAGEConv
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
import torch.nn.functional as F


class Net(torch.nn.Module):  # 针对图进行分类任务
    def __init__(self):
        super(Net, self).__init__()
        #SAGEConv卷积获取邻近结点特征的平均值li 以及自己的特征Xi 即 Xi+1 = W1*Xi+W2*li
        self.conv1 = SAGEConv(embed_dim, 128)
        #TopKPooling池化 ratio=0.8即保留80%的结点
        self.pool1 = TopKPooling(128, ratio=0.8)
        self.conv2 = SAGEConv(128, 128)
        self.pool2 = TopKPooling(128, ratio=0.8)
        self.conv3 = SAGEConv(128, 128)
        self.pool3 = TopKPooling(128, ratio=0.8)
        #df.item_id.max()是item_id列中最大值 加10是为了能够应对新的项
        #torch.nn.Embedding给一个编号，嵌入层就能返回这个编号对应的嵌入向量，嵌入向量反映了各个编号代表的符号之间的语义关系
        self.item_embedding = torch.nn.Embedding(df.item_id.max(), embedding_dim=embed_dim)
        self.lin1 = torch.nn.Linear(128, 128)
        self.lin2 = torch.nn.Linear(128, 64)
        self.lin3 = torch.nn.Linear(64, 1)
        #归一化
        self.bn1 = torch.nn.BatchNorm1d(128)
        self.bn2 = torch.nn.BatchNorm1d(64)
        self.act1 = torch.nn.ReLU()
        self.act2 = torch.nn.ReLU()

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch  # x:n*1,n为图中结点个数其中每个图里点的个数是不同的
        print("data.batch",batch)
        x = self.item_embedding(x)  # n*1*128 特征编码后的结果
        #squeeze降维 原来维度为n*1*128 对应下标0 1 2 squeeze(1)即把下标为1的维度降维变成n*128
        x = x.squeeze(1)  # n*128
        x = F.relu(self.conv1(x, edge_index))  # n*128
        x, edge_index, _, batch, _, _ = self.pool1(x, edge_index, None, batch)  # pool之后得到 n*0.8个点
        print('self.pool1',batch)
        #得到整张图的结点并求平均赋值给x1
        x1 = gap(x, batch)
        x = F.relu(self.conv2(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool2(x, edge_index, None, batch)
        print('pool2',batch)
        x2 = gap(x, batch)
        x = F.relu(self.conv3(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool3(x, edge_index, None, batch)
        x3 = gap(x, batch)
        print('x3',x3.shape)# batch * 256
        x = x1 + x2 + x3  # 获取不同尺度的全局特征

        x = self.lin1(x)
        x = self.act1(x)
        x = self.lin2(x)
        x = self.act2(x)
        #50%的概率关闭神经元 训练时启用 验证时不启用
        x = F.dropout(x, p=0.5, training=self.training)

        x = torch.sigmoid(self.lin3(x)).squeeze(1)  # batch个结果
        # print('sigmoid',x.shape)
        return x


from torch_geometric.loader import DataLoader


def train():
    model.train()

    loss_all = 0
    for data in train_loader:
        data = data
        # print('data',data)
        optimizer.zero_grad()
        output = model(data)
        label = data.y
        loss = crit(output, label)
        loss.backward()
        loss_all += data.num_graphs * loss.item()
        optimizer.step()
    return loss_all / len(dataset)


model = Net()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
crit = torch.nn.BCELoss()
train_loader = DataLoader(dataset, batch_size=64)
for epoch in range(10):
    print('epoch:', epoch)
    loss = train()
    print(loss)

def evalute(loader,model):
    model.eval()

    prediction = []
    labels = []

    with torch.no_grad():
        for data in loader:
            data = data#.to(device)
            pred = model(data)#.detach().cpu().numpy()

            label = data.y#.detach().cpu().numpy()
            prediction.append(pred)
            labels.append(label)
    prediction =  np.hstack(prediction)
    labels = np.hstack(labels)

    return roc_auc_score(labels,prediction)
for epoch in range(1):
    roc_auc_score = evalute(dataset,model)
    print('roc_auc_score',roc_auc_score)
